from dotenv import load_dotenv
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, END

from outline_gen_agents.uilts.outlines_nodes import (
    agent1_requirements, node_macro, node_frontier,
    node_performance, node_other, route_paradigm
)
from outline_gen_agents.uilts.outlines_state import OutLinesState


def build_graph() -> StateGraph:
    g = StateGraph(OutLinesState)
    g.add_node("agent1_requirements", agent1_requirements)
    # 新增四个节点
    g.add_node("宏观态势", node_macro)
    g.add_node("前沿监测", node_frontier)
    g.add_node("绩效评估", node_performance)
    g.add_node("其他", node_other)

    g.set_entry_point("agent1_requirements")

    # 根据 analysis_paradigm 条件路由到对应节点
    g.add_conditional_edges(
        "agent1_requirements",
        route_paradigm,
        {
            "宏观态势": "宏观态势",
            "前沿监测": "前沿监测",
            "绩效评估": "绩效评估",
            "其他": "其他",
        },
    )
    # 每个节点结束后直接终止
    g.add_edge("宏观态势", END)
    g.add_edge("前沿监测", END)
    g.add_edge("绩效评估", END)
    g.add_edge("其他", END)
    return g


def run_agent(user_input: str, provider: str = "deepseek") -> OutLinesState:
    load_dotenv()
    initial = {
        "user_input": user_input,
        "provider": provider
    }
    checkpointer = InMemorySaver()
    graph = build_graph().compile(checkpointer=checkpointer)
    config = {"configurable": {"thread_id": "1"}}
    state = graph.invoke(initial, config=config)
    return state


if __name__ == "__main__":
    run_agent("面向人工智能领域，分析最近几年的前沿进展并形成报告", "deepseek")
